# 测试单张图片的代码
import json
import ctypes
import os
import shutil
import sys
import cv2
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import torch
import court  # 导入场地检测模块

# 从原始文件导入YoLov5TRT类
from yolov5_tensorrt_venueVersion import YoLov5TRT, Colors, plot_one_box, categories

# 初始化全局变量
person_x = 0
person_y = 0
G_x, G_y = 0, 0
lines = []
left_line = []
Right_line = []
top_line = []
bottom_line = []


def test_single_image(img_path):
    """测试单张图片的检测和场地识别"""
    global person_x, person_y, G_x, G_y, lines, left_line, Right_line, top_line, bottom_line

    # 加载自定义插件和引擎
    PLUGIN_LIBRARY = "build/libmyplugins.so"
    engine_file_path = "build/yolov5s.engine"

    if len(sys.argv) > 1:
        img_path = sys.argv[1]
    if len(sys.argv) > 2:
        engine_file_path = sys.argv[2]
    if len(sys.argv) > 3:
        PLUGIN_LIBRARY = sys.argv[3]

    # 创建输出目录
    output_dir = os.getcwd() + "/output/"
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    # 加载TensorRT引擎
    ctypes.CDLL(PLUGIN_LIBRARY)
    yolov5_wrapper = YoLov5TRT(engine_file_path)

    try:
        # 读取图片
        img = cv2.imread(img_path)
        if img is None:
            print(f"无法读取图片: {img_path}")
            return

        # 调整图片大小
        img = cv2.resize(img, (800, 600))
        image_copy = img.copy()

        # 执行目标检测
        fake_result, use_time = yolov5_wrapper.infer(img)
        print(f"检测用时: {use_time * 1000:.2f}ms")
        print(f"检测结果: {fake_result}")

        # 场地检测
        lines = court.Get_Lines(image_copy)

        # 判断是竖线还是横线
        court_exist, h_lines, s_lines = court.horizontal_vertical(lines)

        if court_exist:
            # 从大到小排序线
            court.sort(h_lines)
            # 清除多余的线
            court.clear_duplicate_lines(h_lines)

            lines = h_lines + s_lines

            print("检测到的线条:")
            print(lines)

            # 寻找球场线
            left_line, Right_line, top_line, bottom_line = court.find_course(lines)
            print("场地线:")
            print(f"左线: {left_line}")
            print(f"右线: {Right_line}")
            print(f"顶线: {top_line}")
            print(f"底线: {bottom_line}")

            # 绘制球场线
            img = court.draw_course(img, left_line, Right_line, top_line, bottom_line)

            # 如果检测到人员且场地存在，则计算场地坐标
            if person_x > 0 and person_y > 0:
                # 计算场地坐标
                if len(left_line) and len(Right_line):
                    x1, y1, x2, y2, flag = left_line
                    left_x = (person_y - y1) * (x2 - x1) / (y2 - y1) + x1

                    x1, y1, x2, y2, flag = Right_line
                    Right_x = (person_y - y1) * (x2 - x1) / (y2 - y1) + x1

                    pixel_to_CM = 610 / abs(Right_x - left_x)  # 一个像素点对应的CM
                    G_x = (person_x - left_x) * pixel_to_CM  # X相对坐标

                    pixel_to_CM = (410 / abs(bottom_line[1] - top_line[1]))  # 一个像素点对应的CM
                    G_y = (((lines[0][0][1] + lines[0][0][3])) / 2 - person_y) * pixel_to_CM  # Y相对坐标

                    G_x = int(G_x)
                    G_y = int(G_y)

                    print(f"人员像素坐标: ({person_x}, {person_y})")
                    print(f"场地坐标: ({G_x}, {G_y}) cm")

                    # 在图像上标注人员的场地坐标
                    cv2.putText(img, f"Court Pos: ({G_x},{G_y})cm",
                                (person_x - 50, person_y - 20),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)

                    # 在人员位置画一个圆点标记
                    cv2.circle(img, (person_x, person_y), 5, (0, 255, 255), -1)
        else:
            print("未检测到完整场地")

        # 保存结果图像
        parent, filename = os.path.split(img_path)
        cv2.imwrite(output_dir + filename, img)
        print(f"结果已保存至: {output_dir + filename}")

        # 显示结果图像
        cv2.imshow("Detection Result", img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    finally:
        # 释放资源
        yolov5_wrapper.destroy()


if __name__ == "__main__":
    # 指定要测试的图片路径
    img_path = "image/test1.jpg"  # 替换为你的图片路径
    test_single_image(img_path)
